import os
from tqdm import tqdm
import json 
import numpy as np 
from shapely import geometry,affinity 
def rotatePoints(points,center,theta):
    '''
    theta角度为360度制
    '''
    # assert 360>=theta and theta >=0 
    newPoints = points.copy()
    radTheta = theta*np.math.pi/180
    newPoints[:,0] = (points[:,0]-center[0])*np.math.cos(radTheta) - (points[:,1]-center[1])*np.math.sin(radTheta)+center[0]
    newPoints[:,1] = (points[:,0]-center[0])*np.math.sin(radTheta) + (points[:,1]-center[1])*np.math.cos(radTheta)+center[1]
    return newPoints

def getRematchPoint(pointsOfOrigin,pointsOfRect):
    '''
    这个函数用来重匹配各个点
    '''
    assert pointsOfOrigin.shape == pointsOfRect.shape
    pointsOfRect = pointsOfRect.copy()
    pointsHaveMatch = []
    for i,point in enumerate (pointsOfOrigin):
        lengthArray = np.linalg.norm(point-pointsOfRect,axis=-1)
        minIndex = np.argmin(lengthArray)
        pointsHaveMatch.append(pointsOfRect[minIndex])
        np.delete(pointsOfRect,minIndex,axis=0)
    return np.array(pointsHaveMatch)

def getVectorsAngle(vector):
    '''
    得到向量的绝对角度
    '''
    radTheta = np.math.atan2(vector[-1],vector[0])
    degTheta = np.math.degrees(radTheta)
    return degTheta

    

def getYoloLabel(pointsOfOrigin,mode='360'):
    '''
    得到yolo格式的标签:  (x,y),(w,h),θ
    得到dota格式的标签: (x1,y1),...,(x4,y4)
    '''
    mpt = geometry.MultiPoint(pointsOfOrigin)
    minRectPolygen = mpt.minimum_rotated_rectangle
    minRectPoints = np.array(minRectPolygen.boundary)[:4]
    # print(minRectPolygen.boundary)
    # print(pointsOfOrigin.shape,minRectPoints.shape)
    rematchPoints = getRematchPoint(pointsOfOrigin,minRectPoints)
    w = np.linalg.norm(rematchPoints[0]-rematchPoints[1])
    h = np.linalg.norm(rematchPoints[1]-rematchPoints[2])
    theta = getVectorsAngle(rematchPoints[0]-rematchPoints[-1])
    (x,y) = np.array(minRectPolygen.centroid)
    if mode == '180':
        theta =  theta + 180 if theta < 0 else theta
    if mode == '360':
        theta = 360 + theta if theta < 0 else theta
    return {'YoloLabel':((x,y),(w,h),theta),'DotaLabel':rematchPoints}
    
def getClasses(labelPath):
    # imgPath = r'D:\dataset\DOTA\train\images'
    # xIndex = [0,2,4,6]
    # yIndex = [1,3,5,7]
    classes = []
    labels = []
    nameList = os.listdir(labelPath)
    nameList.sort()
    for name in tqdm(nameList):
        location =  labelPath + '/' + name
        # nameOnly = name.split('.')[0]
        with open(location,'r') as txtFile:
            lines = txtFile.read().splitlines()
            for j,line in enumerate(lines):
                if len(line)<=24:
                    continue
                newLine = np.array(line.split(' '))
                points = newLine[:8].reshape(-1,2).astype(np.int64)
                label = getYoloLabel(points)
                yoloLabel = label['YoloLabel']
                # (x,y),(w,h),theta = yoloLabel
                # dotaLabel = label['DotaLabel']
                labels.append(yoloLabel)
                singleClass = newLine[8]
                if singleClass not in classes:
                    classes.append(singleClass)
    return{'classes':classes,'labels':labels}

if __name__ == '__main__':
    txtPath = r'D:\dataset\DOTA\train\labelTxt-v1.0\labelTxt'
    jsonPath = './data/airplane/annotations/train.json'
    with open(jsonPath,'r') as file:
        labelDict = json.load(file)
        print(labelDict.keys())
        print(labelDict['images'][0].keys())
        print(labelDict['categories'][0].keys())
        print(labelDict['annotations'][0].keys())
    labelData = getClasses(txtPath)
    # print(a)